#include "libmri.h"

#include <algorithm>
#include <iostream>
using std::min;
using std::max;
using std::cout;
using std::endl;

__global__ void
cukreduce3d (cplxf *KKt, cplxf *K, int ks1, int ks2, int ks3, int n, int m)
{
	int kk1 = 2*ks1-1, kk2 = 2*ks2-1, kk3 = 2*ks3-1;

	for (int bid = blockIdx.x; bid < n*n; bid += gridDim.x)
	{
		int i = bid / n;
		int j = bid - n*i;

		for (int k = 0; k < m; ++k)
		{
			cplxf *KKtij = &KKt [(i + n*j)*kk1*kk2*kk3],
			        *Kik =   &K [(i + n*k)*ks1*ks2*ks3],
			        *Kjk =   &K [(j + n*k)*ks1*ks2*ks3];

			for (int tid = threadIdx.x; tid < kk1*kk2*kk3; tid += blockDim.x)
			{
				int i1 = tid;
				int i3 = i1 / (kk1*kk2); i1 -= i3*kk1*kk2;
				int i2 = i1 / kk1;       i1 -= i2*kk1;

				int j3_0 = max(0, i3-ks3+1), j3_n = min(i3+1, ks3);
				int j2_0 = max(0, i2-ks2+1), j2_n = min(i2+1, ks2);
				int j1_0 = max(0, i1-ks1+1), j1_n = min(i1+1, ks1);

				cplxf z = KKtij [i1 + kk1*(i2 + kk2*i3)];

				for (int j3 = j3_0; j3 < j3_n; ++j3)
				  for (int j2 = j2_0; j2 < j2_n; ++j2)
					for (int j1 = j1_0; j1 < j1_n; ++j1)
					{
						z += Kik [(i1-j1) + ks1*
							     ((i2-j2) + ks2*
							      (i3-j3))]
							 * conj (Kjk [(ks1-1-j1) + ks1*
							             ((ks2-1-j2) + ks2*
							              (ks3-1-j3))]);
					}

				KKtij [i1 + kk1*(i2 + kk2*i3)] = z; 
			}
		}
	}
}

void cukreduce3d (cplxf *_KKt, cplxf *_K, int *ksize, int n, int m)
{
	int ks1 = ksize[0], ks2 = ksize[1], ks3 = ksize[2];
	int kk1 = 2*ks1-1, kk2 = 2*ks2-1, kk3 = 2*ks3-1;

	cudaDeviceProp prop;
	int dev;

	cuda(GetDevice(&dev));
	cuda(GetDeviceProperties(&prop,dev));


	cplxf *KKt; cuda(Malloc(&KKt, kk1*kk2*kk3*n*n*sizeof(cplxf)));
	cplxf *K; cuda(Malloc (&K, ks1*ks2*ks3*n*m*sizeof(cplxf)));

	/*
	double ts = timestamp();
	*/

	cuda(Memcpy (K, _K, ks1*ks2*ks3*n*m*sizeof(cplxf), cudaMemcpyHostToDevice));
	cuda(Memset (KKt, 0, kk1*kk2*kk3*n*n*sizeof(cplxf)));

	/*
	cuda(ThreadSynchronize());
	cout << "memcpy() to device: " << (timestamp()-ts) << " s" << endl;
	ts = timestamp();
	*/

	cukreduce3d <<< 2*prop.multiProcessorCount, 64 >>>
	              ( KKt, K, ks1, ks2, ks3, n, m);

	cuda(ThreadSynchronize());

	/*
	cout << "kernel call: " << (timestamp()-ts) << " s" << endl;
	ts = timestamp();
	*/

	cuda(Memcpy (_KKt, KKt, kk1*kk2*kk3*n*n*sizeof(cplxf), cudaMemcpyDeviceToHost));

	/*
	cuda(ThreadSynchronize());
	cout << "memcpy() to host: " << (timestamp()-ts) << " s" << endl;
	*/

	cuda (Free (KKt));
	cuda (Free (K));
}

